import cv2
import numpy as np
from scipy.interpolate import CubicSpline


def process_image(model, image_path):
    """
    读取图像并进行目标检测，智能识别主体物并绘制轮廓

    参数:
        model (YOLO): 加载的 YOLO 分割模型
        image_path (str): 图像文件路径

    返回:
        np.ndarray: 带有检测结果的图像
        list: 检测到的物品列表
    """
    # 初始化检测到的物品列表
    detected_items = []

    # 读取输入图像
    image = cv2.imread(image_path)
    if image is None:
        print("无法读取图像")
        return None, detected_items

    # 进行目标检测与分割
    results = model(image, conf=0.8, verbose=False)  # 设置置信度阈值并关闭详细输出

    # 获取图像中心
    h, w, _ = image.shape
    center = np.array([w // 2, h // 2])

    # 存储检测到的物体信息
    detected_objects = []

    # 解析检测结果
    for result in results:
        masks = result.masks  # 获取分割掩码
        boxes = result.boxes  # 获取边界框

        if masks is None:
            continue

        for i, box in enumerate(boxes):
            # 获取边界框和类别信息
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            confidence = float(box.conf[0])
            class_id = int(box.cls[0])
            item_name = model.names[class_id]
            detected_items.append(item_name)  # 添加检测到的物品

            # 获取对应的掩码并提取轮廓
            mask = masks.data[i].cpu().numpy()
            mask = cv2.resize(mask, (w, h))
            _, mask = cv2.threshold(mask, 0.5, 255, cv2.THRESH_BINARY)
            mask = mask.astype(np.uint8)

            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            # 计算物体面积（像素数量）
            area = cv2.contourArea(np.vstack(contours)) if contours else 0

            # 计算物体中心（使用轮廓质心）
            if contours:
                M = cv2.moments(np.vstack(contours))
                cx = int(M['m10'] / M['m00']) if M['m00'] != 0 else (x1 + x2) // 2
                cy = int(M['m01'] / M['m00']) if M['m00'] != 0 else (y1 + y2) // 2
                obj_center = np.array([cx, cy])
            else:
                obj_center = np.array([(x1 + x2) // 2, (y1 + y2) // 2])

            # 计算到图像中心的归一化距离
            max_distance = np.sqrt(w ** 2 + h ** 2)
            distance = np.linalg.norm(obj_center - center) / max_distance

            # 计算大小因子（归一化面积）
            max_area = w * h
            size_factor = 1 - (area / max_area)

            # 综合评分（距离权重50%，大小权重30%，置信度权重20%）
            score = (0.5 * (1 - distance)) + (0.3 * (1 - size_factor)) + (0.2 * confidence)

            detected_objects.append({
                'contours': contours,
                'confidence': confidence,
                'class_id': class_id,
                'score': score,
                'label': model.names[class_id],
                'center': obj_center
            })

    # 如果检测到物体，选择评分最高的作为主体物
    if detected_objects:
        detected_objects.sort(key=lambda x: x['score'], reverse=True)
        main_obj = detected_objects[0]
        color = (0, 255, 0)  # 绿色

        # 创建遮罩层
        overlay = image.copy()

        # 优化1：合并所有轮廓点
        all_points = []
        for contour in main_obj['contours']:
            all_points.extend(contour)

        # 优化2：直接使用原始轮廓点并裁剪到边界
        if all_points:
            clipped_points = []
            for point in all_points:
                x, y = point[0]
                # 确保点坐标在图像范围内
                x = max(0, min(x, w - 1))
                y = max(0, min(y, h - 1))
                clipped_points.append([[x, y]])

            # 使用原始轮廓点而非凸包
            contour_array = np.array(clipped_points, dtype=np.int32)
            cv2.fillPoly(overlay, [contour_array], color)
        else:
            # 点数不足时使用原始轮廓并裁剪到边界
            for contour in main_obj['contours']:
                contour_clipped = []
                for point in contour:
                    x, y = point[0]
                    # 确保点坐标在图像范围内
                    x = max(0, min(x, w - 1))
                    y = max(0, min(y, h - 1))
                    contour_clipped.append([[x, y]])

                cv2.fillPoly(overlay, [np.array(contour_clipped, dtype=np.int32)], color)

        # 设置透明度并叠加到原图
        alpha = 0.4  # 40%不透明
        image = cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0)

        # 绘制轮廓描边（智能曲线检测）
        for contour in main_obj['contours']:
            if len(contour) >= 4:  # 需要至少4个点
                # 提取轮廓点并应用高斯滤波
                points = contour[:, 0, :].astype(np.float32)

                # 修复高斯滤波形状不匹配问题
                x_coords = points[:, 0].reshape(-1)
                y_coords = points[:, 1].reshape(-1)

                # 应用高斯滤波
                x_smooth = cv2.GaussianBlur(x_coords, (5, 1), 0)
                y_smooth = cv2.GaussianBlur(y_coords, (5, 1), 0)

                # 组合处理后的坐标
                points = np.column_stack((x_smooth, y_smooth))
                smooth_contour = points.reshape(-1, 1, 2).astype(np.int32)

                # 计算轮廓复杂度（曲线度）
                orig_length = cv2.arcLength(smooth_contour, True)
                epsilon = 0.005 * orig_length
                approx_contour = cv2.approxPolyDP(smooth_contour, epsilon, True)
                approx_length = cv2.arcLength(approx_contour, True)

                # 计算复杂度 = 原始长度 / 逼近长度
                complexity = orig_length / approx_length if approx_length > 0 else 1.0

                # 智能绘制策略
                if complexity < 1.05:  # 接近直线
                    # 直接绘制简化后的直线轮廓
                    cv2.polylines(image, [approx_contour], isClosed=True,
                                  color=color, thickness=25, lineType=cv2.LINE_AA)
                else:  # 曲线
                    try:
                        # 参数化
                        t = np.linspace(0, 1, len(points))
                        t_new = np.linspace(0, 1, 100)  # 100个点用于插值

                        # 三次样条插值
                        cs_x = CubicSpline(t, points[:, 0])
                        cs_y = CubicSpline(t, points[:, 1])

                        x_new = cs_x(t_new).astype(int)
                        y_new = cs_y(t_new).astype(int)

                        curve_contour = np.stack([x_new, y_new], axis=1)
                        curve_contour = curve_contour.reshape(-1, 1, 2)

                        # 绘制平滑曲线
                        cv2.polylines(image, [curve_contour], isClosed=True,
                                      color=color, thickness=25, lineType=cv2.LINE_AA)
                    except Exception:
                        # 回退到原始轮廓
                        cv2.polylines(image, [smooth_contour], isClosed=True,
                                      color=color, thickness=25, lineType=cv2.LINE_AA)
            else:
                # 点数不足时直接绘制
                cv2.polylines(image, [contour], isClosed=True,
                              color=color, thickness=25, lineType=cv2.LINE_AA)

    # 绘制标签
    if main_obj['contours']:
        x, y, _, _ = cv2.boundingRect(main_obj['contours'][0])
        main_object = main_obj['label']
        accuracy_rate = f"{main_obj['confidence']:.2f}"
        label = f"{main_object} {accuracy_rate}"
        # print("检测到主体物: ", main_object)
        # print("置信度: ", accuracy_rate)
        # cv2.putText(image, label, (x, y - 30),
        #             cv2.FONT_HERSHEY_SIMPLEX, 2, color, 4, cv2.LINE_AA)

    return image, detected_items, main_object, accuracy_rate
